﻿/*
I certainly find the enumerable methods here useful, but there's more to add to
this and already it's mostly the non-enumerable aspects that are going to grow.
Should probably create a completely separate project for this. 
*/

using System.Threading;

namespace MoreLinq
{
    using System;
    using System.Collections;
    using System.Collections.Generic;
    internal sealed class LockUse : IDisposable
    {
        public enum State
        {
            None,
            Shared,
            Exclusive//Not using upgradable here.
        }
        private readonly ReaderWriterLock _lock;
        private State _state;
        public LockUse(ReaderWriterLock rwl, State state)
        {
            _lock = rwl;
            switch(_state = state)
            {
                case State.Shared:
                    rwl.AcquireReaderLock(-1);
                    break;
                case State.Exclusive:
                    rwl.AcquireWriterLock(-1);
                    break;
            }
        }
        public void ClearLock()
        {
            _lock.ReleaseLock();
            _state = State.None;
        }
        public void EnterShared()
        {
            switch(_state)
            {
                case State.Exclusive:
                    _lock.ReleaseLock();
                    goto case State.None; //explicit fallthrough
                case State.None:
                    _lock.AcquireReaderLock(-1);
                    break;
            }
            _state = State.Shared;
        }
        public void EnterExclusive()
        {
            switch(_state)
            {
                case State.Shared:
                    _lock.ReleaseLock();
                    goto case State.None; //explitic fallthrough
                case State.None:
                    _lock.AcquireWriterLock(-1);
                    break;
            }
            _state = State.Exclusive;
        }
        public void Dispose()
        {
            ClearLock();
        }
    }
    internal static class InternHelper
    {
        private static readonly int[] Primes = new int[]
        {
            7, 11, 17, 23, 29, 37, 47, 59, 71, 89, 107, 131, 163, 197, 239, 293, 353, 431, 521,
            631, 761, 919, 1103, 1327, 1597, 1931, 2333, 2801, 3371, 4049, 4861, 5839, 7013, 8419,
            10103, 12143, 14591, 17519, 21023, 25229, 30293, 36353, 43627, 52361, 62851, 75431, 90523,
            108631, 130363, 156437, 187751, 225307, 270371, 324449, 389357, 467237, 560689, 672827,
            807403, 968897, 1162687, 1395263, 1674319, 2009191, 2411033, 2893249, 3471899, 4166287,
            4999559, 5999471, 7199369, 10000019, 15000017, 20000003, 32452843, 49979687, 67867967,
            86028121, 104395301, 160481183, 217645177, 256203161, 314606869, 413158511, 512927357,
            593441843, 654188383, 735632791, 858599503, 899809363, 982451653
        };
        public static int NextSize(int targetSize)
        {
            if(targetSize < 7)
                return 7;
            else if(targetSize <= 982451653)
            {
                int idx = Array.BinarySearch(Primes, targetSize);
                if(idx >= 0)
                    return targetSize;
                else
                    return Primes[~idx];
            }
            else if(targetSize >= 0x7FFFFFFF)
                return 0x7FFFFFFF;
            else
                return targetSize; //forget about primality and just use the min as is.
        }
    }
    /// <summary>
    /// 
    /// </summary>
    public abstract class InternPool<T> where T : class
    {
        private readonly IEqualityComparer<T> _comparer;
        /// <summary>
        /// 
        /// </summary>
        public IEqualityComparer<T> Comparer
        {
            get { return _comparer; }
        }
        internal InternPool(IEqualityComparer<T> comparer)
        {
            _comparer = comparer ?? EqualityComparer<T>.Default;
        }
        internal int GetHashCode(T value)
        {
            return _comparer.GetHashCode(value);
        }
        internal bool Equals(T x, T y)
        {
            return _comparer.Equals(x, y);
        }
        internal static int IndexFromHash(int hash, int length)
        {
            return (hash & 0x7FFFFFFF) % length;
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="value"></param>
        /// <returns></returns>
        public abstract T Intern(T value);
        public abstract T IsInterned(T value);
        /// <summary>
        /// 
        /// </summary>
        /// <param name="source"></param>
        /// <returns></returns>
        public IEnumerable<T> Intern(IEnumerable<T> source)
        {
            if (source == null) throw new ArgumentNullException("source");
            return new InterningEnumerable(this, source);
        }
        private class InterningEnumerable : IEnumerable<T>
        {
            private readonly InternPool<T> _pool;
            private readonly IEnumerable<T> _source;
            public InterningEnumerable(InternPool<T> pool, IEnumerable<T> source)
            {
                _pool = pool;
                _source = source;
            }
            public IEnumerator<T> GetEnumerator()
            {
                return new InterningEnumerator(_pool, _source);
            }
            IEnumerator IEnumerable.GetEnumerator()
            {
                return GetEnumerator();
            }
        }
        private class InterningEnumerator : IEnumerator<T>
        {
            private readonly InternPool<T> _pool;
            private readonly IEnumerator<T> _sourceEn;
            private T _current;
            public InterningEnumerator(InternPool<T> pool, IEnumerable<T> source)
            {
                _pool = pool;
                _sourceEn = source.GetEnumerator();
            }
            public T Current
            {
                get { return _current; }
            }
            object IEnumerator.Current
            {
                get { return Current; }
            }
            public void Dispose()
            {
                _sourceEn.Dispose();
            }
            public bool MoveNext()
            {
                if(_sourceEn.MoveNext())
                {
                    _current = _pool.Intern(_sourceEn.Current);
                    return true;
                }
                return false;
            }
            public void Reset()
            {
                _sourceEn.Reset();
            }
        }
    }
    /// <summary>
    /// 
    /// </summary>
    public sealed class StrongInternPool<T> : InternPool<T> where T : class
    {
        private struct Slot
        {
            public int Hash;
            public int Next;
            public T Value;
        }
        private int[] _indices;
        private Slot[] _slots;
        private int _freeList = -1;
        private int _addingIndex = 0;
        /// <summary>
        /// 
        /// </summary>
        /// <param name="comparer"></param>
        /// <param name="capacity"></param>
        public StrongInternPool(IEqualityComparer<T> comparer, int capacity)
            :base(comparer)
        {
            capacity = InternHelper.NextSize(capacity);
        	_indices = new int[capacity];
        	_slots = new Slot[capacity];
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="capacity"></param>
        public StrongInternPool(int capacity)
            :this(EqualityComparer<T>.Default, capacity){}
        /// <summary>
        /// 
        /// </summary>
        /// <param name="comparer"></param>
        public StrongInternPool(IEqualityComparer<T> comparer)
            :this(comparer, 7){}
        /// <summary>
        /// 
        /// </summary>
        public StrongInternPool()
            :this(EqualityComparer<T>.Default, 7){}
        /// <summary>
        /// 
        /// </summary>
        /// <param name="value"></param>
        /// <returns></returns>
        public override T IsInterned(T value)
        {
            if(ReferenceEquals(value, null))
                return null;
            int hash = GetHashCode(value);
            int idx = IndexFromHash(hash, _indices.Length);
        	for(int i = _indices[idx] - 1; i >= 0; i = _slots[i].Next)
        	    if(_slots[i].Hash == hash)
        	    {
        	        T val = _slots[i].Value;
            		if(Equals(val, value))
            		    return val;
        	    }
        	return null;
        }
        public override T Intern(T value)
        {
            if(ReferenceEquals(value, null))
                return null;
            int hash = GetHashCode(value);
            int idx = IndexFromHash(hash, _indices.Length);
        	for(int i = _indices[idx] - 1; i >= 0; i = _slots[i].Next)
        	    if(_slots[i].Hash == hash)
        	    {
        	        T val = _slots[i].Value;
            		if(Equals(val, value))
            		    return val;
        	    }
        	int addIdx;
        	if(_freeList >= 0)
        	{
        		addIdx = _freeList;
        		_freeList = _slots[addIdx].Next;
        	}
        	else
        	{
        		if(_addingIndex == _slots.Length)
        		{
        			IncreaseCapacity();
        			idx = IndexFromHash(hash, _indices.Length);
        		}
        		addIdx = _addingIndex;
        		++_addingIndex;
        	}
        	_slots[addIdx].Hash = hash;
        	_slots[addIdx].Value = value;
        	_slots[addIdx].Next = _indices[idx] - 1;
        	_indices[idx] = addIdx + 1;
        	return value;
        }
        private void IncreaseCapacity()
        {
        	int newSize = InternHelper.NextSize(2 * _addingIndex);
        	if(newSize <= _addingIndex)
        	    throw new ArgumentException("Intern Pool Overflow");
        	Slot[] newSlots = new Slot[newSize];
    		Array.Copy(_slots, 0, newSlots, 0, _addingIndex);
        	int[] _newIndices = new int[newSize];
        	for(int i = 0; i < _addingIndex; ++i)
        	{
        	    int idx = IndexFromHash(newSlots[i].Hash, newSize);
        		newSlots[i].Next = _newIndices[idx] - 1;
        		_newIndices[idx] = i + 1;
        	}
        	_slots = newSlots;
        	_indices = _newIndices;
        }
    }
    public sealed class StripedInternPool<T> : InternPool<T> where T : class
    {
        private class Node
        {
            public int Hash;
            public T Value;
            public Node Next;
        }
        private Node[] _slots;
        private readonly ReaderWriterLock _lock = new ReaderWriterLock();
        public StripedInternPool(IEqualityComparer<T> comparer, int capacity)
            :base(comparer)
        {
            _slots = new Node[InternHelper.NextSize((int)Math.Sqrt(capacity))];
        }
        public StripedInternPool()
            :this(EqualityComparer<T>.Default, 7){}
        public override T IsInterned(T value)
        {
            if(ReferenceEquals(value, null))
                return null;
            int hash = GetHashCode(value);
            int idx = IndexFromHash(hash, _slots.Length);
            for(Node node = _slots[idx]; node != null; node = node.Next)
                if(node.Hash == hash)
                {
                    T val = node.Value;
                    if(Equals(val, value))
                        return val;
                }
            return null;
        }
        public override T Intern(T value)
        {
            return Intern(value, false);
        }
        private T Intern(T value, bool noGrow)
        {
            if(ReferenceEquals(value, null))
                return null;
            int hash = GetHashCode(value);
            using(LockUse lockUse = new LockUse(_lock, LockUse.State.Shared))
            {
                int len = _slots.Length;
                int idx = IndexFromHash(hash, len);
                Node node = _slots[idx];
                if(node == null)
                {
                    node = Interlocked.CompareExchange(ref _slots[idx], new Node{Hash = hash, Value = value}, null);
                    if(node == null)//CAS worked
                        return value;
                }
                int loopCount = 0;
                for(;;)
                {
                    if(node.Hash == hash && Equals(node.Value, value))
                        return node.Value;
                    Node next = node.Next;
                    if(next != null)
                        node = next;
                    else if(noGrow || loopCount > len)
                    {
                        lockUse.EnterExclusive();
                        T ret = Grow(value, hash, loopCount, _slots);
                        if(ret != null)
                            return ret;
                        break;
                    }
                    else
                    {
                        node = Interlocked.CompareExchange(ref node.Next, new Node{Hash = hash, Value = value}, null);
                        if(node == null)//CAS worked
                            return value;
                    }
                    ++loopCount;
                }
            }
            return Intern(value, true);
        }
        private IEnumerable<Node> AllNodes(Node[] slots, T addValue, int hash)
        {
            for(var i = 0; i != slots.Length; ++i)
                for(Node node = slots[i]; node != null; node = node.Next)
                    yield return node;
            yield return new Node{Hash = hash, Value = addValue};
        }
        private T Grow(T value, int hash, int loopCount, Node[] curSlots)
        {
            if(curSlots != _slots)
                return null;
            int newLen = InternHelper.NextSize(loopCount * 2);
            Node[] newSlots = new Node[newLen];
            foreach(Node node in AllNodes(curSlots, value, hash))
            {
                hash = node.Hash;
                Node toAdd = new Node{Hash = hash, Value = node.Value};
                int idx = IndexFromHash(hash, newLen);
                Node addAt = newSlots[idx];
                if(addAt == null)
                    newSlots[idx] = toAdd;
                else
                {
                    Node next = addAt.Next;
                    while(next != null)
                        next = (addAt = next).Next;
                    addAt.Next = toAdd;
                }
            }
            _slots = newSlots;
            return value;
        }
    }
    public sealed class StripedWeakInternPool<T> : InternPool<T> where T : class
    {
        private class Node
        {
            public int Hash;
            public WeakReference Reference;
            public Node Next;
        }
        private Node[] _slots;
        private int _deadClearVersion;
        private readonly ReaderWriterLock _lock = new ReaderWriterLock();
        public StripedWeakInternPool(IEqualityComparer<T> comparer, int capacity)
            :base(comparer)
        {
            _slots = new Node[InternHelper.NextSize((int)Math.Sqrt(capacity))];
        }
        public StripedWeakInternPool()
            :this(EqualityComparer<T>.Default, 7){}
        public override T Intern(T value)
        {
            return Intern(value, false);
        }
        private T Intern(T value, bool noGrow)
        {
            if(ReferenceEquals(value, null))
                return null;
            int hash = GetHashCode(value);
            using(LockUse lockUse = new LockUse(_lock, LockUse.State.Shared))
            {
                int len = _slots.Length;
                int idx = IndexFromHash(hash, len);
                Node node = _slots[idx];
                bool seenDead = false;
                Node yoink = null;
                if(node == null)
                {
                    node = Interlocked.CompareExchange(ref _slots[idx], new Node{Hash = hash, Reference = new WeakReference(value)}, null);
                    if(node == null)//CAS worked
                        return value;
                }
                int loopCount = 0;
                for(;;)
                {
                    if(!seenDead)
                        seenDead = !node.Reference.IsAlive;
                    if(node.Hash == hash)
                    {
                        T val = (T)node.Reference.Target;
                        if(ReferenceEquals(val, null))
                        {
                            if(yoink == null)
                                yoink = node;
                        }
                        else if(Equals(val, value))
                            return val;
                    }
                    Node next = node.Next;
                    if(next != null)
                        node = next;
                    else
                    {
                        if(yoink != null)//reuse a now-dead reference;
                        {
                            WeakReference wr = yoink.Reference;
                            if(
                                !wr.IsAlive //still dead, not replaced by something else
                                &&
                                Interlocked.CompareExchange(ref yoink.Reference, new WeakReference(value), wr) == wr//CAS succeeded
                              )
                                return value;
                        }
                        if(noGrow || loopCount > len)
                        {
                            if(seenDead)
                            {
                                if(!ClearDead(lockUse, _slots, _deadClearVersion))
                                    break;
                            }
                            else
                            {
                                lockUse.EnterExclusive();
                                T ret = Grow(value, hash, loopCount, _slots);
                                if(ret != null)
                                    return ret;
                                break;
                            }
                        }
                        node = Interlocked.CompareExchange(ref node.Next, new Node{Hash = hash, Reference = new WeakReference(value)}, null);
                        if(node == null)//CAS worked
                            return value;
                    }
                    ++loopCount;
                }
            }
            return Intern(value, true);
        }
        public override T IsInterned(T value)
        {
            if(ReferenceEquals(value, null))
                return null;
            int hash = GetHashCode(value);
            using(LockUse lockUse = new LockUse(_lock, LockUse.State.Shared))
            {
                int len = _slots.Length;
                int idx = IndexFromHash(hash, len);
                for(Node node = _slots[idx]; node != null; node = node.Next)
                    if(node.Hash == hash)
                    {
                        T val = (T)node.Reference.Target;
                        if(!ReferenceEquals(val, null) && Equals(val, value))
                            return val;
                    }
                return null;
            }
        }
        private IEnumerable<Node> AllNodes(Node[] slots, T addValue, int hash)
        {
            for(var i = 0; i != slots.Length; ++i)
                for(Node node = slots[i]; node != null; node = node.Next)
                    if(node.Reference.IsAlive)
                        yield return node;
            yield return new Node{Hash = hash, Reference = new WeakReference(addValue)};
        }
        private T Grow(T value, int hash, int loopCount, Node[] curSlots)
        {
            if(curSlots != _slots)
                return null;
            int newLen = InternHelper.NextSize(loopCount * 2);
            Node[] newSlots = new Node[newLen];
            foreach(Node node in AllNodes(curSlots, value, hash))
            {
                hash = node.Hash;
                Node toAdd = new Node{Hash = hash, Reference = node.Reference};
                int idx = IndexFromHash(hash, newLen);
                Node addAt = newSlots[idx];
                if(addAt == null)
                    newSlots[idx] = toAdd;
                else
                {
                    Node next = addAt.Next;
                    while(next != null)
                        next = (addAt = next).Next;
                    addAt.Next = toAdd;
                }
            }
            _slots = newSlots;
            return value;
        }
        private bool ClearDead(LockUse lockUse, Node[] slots, int deadClearVersion)
        {
            lockUse.EnterExclusive();
            if(_deadClearVersion != deadClearVersion)
                return false;
            for(int i = 0; i != slots.Length; ++i)
            {
                Node node = slots[i];
                while(node != null && !node.Reference.IsAlive)
                    slots[i] = node = node.Next;
                if(node != null)
                {
                    Node next = node.Next;
                    while(next != null)
                    {
                        if(!next.Reference.IsAlive)
                            node.Next = next.Next;
                        else
                            node = next;
                        next = node.Next;
                    }
                }
            }
            lockUse.EnterShared();
            return true;
        }
    }
    internal struct WeakSlot
    {
        public int Hash;
        public int Next;
        public WeakReference Reference;
    }
    /// <summary>
    /// 
    /// </summary>
    public sealed class WeakInternPool<T> : InternPool<T> where T : class
    {
        private int[] _indices;
        private WeakSlot[] _slots;
        private int _freeList = -1;
        private int _addingIndex = 0;
        /// <summary>
        /// 
        /// </summary>
        /// <param name="comparer"></param>
        /// <param name="capacity"></param>
        public WeakInternPool(IEqualityComparer<T> comparer, int capacity)
            :base(comparer)
        {
            capacity = InternHelper.NextSize(capacity);
        	_indices = new int[capacity];
        	_slots = new WeakSlot[capacity];
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="capacity"></param>
        public WeakInternPool(int capacity)
            :this(EqualityComparer<T>.Default, capacity){}
        /// <summary>
        /// 
        /// </summary>
        /// <param name="comparer"></param>
        public WeakInternPool(IEqualityComparer<T> comparer)
            :this(comparer, 7){}
        /// <summary>
        /// 
        /// </summary>
        public WeakInternPool()
            :this(EqualityComparer<T>.Default, 7){}
        /// <summary>
        /// 
        /// </summary>
        /// <param name="value"></param>
        /// <returns></returns>
        public override T IsInterned(T value)
        {
            if(!ReferenceEquals(value, null))
            {
                int hash = GetHashCode(value);
                int idx = IndexFromHash(hash, _indices.Length);
                for(int i = _indices[idx] - 1; i >= 0; i = _slots[i].Next)
            	    if(_slots[i].Hash == hash)
            	    {
            	        T target = (T)_slots[i].Reference.Target;
            	        if(!ReferenceEquals(target, null) && Equals(target, value))
                		    return target;
            	    }
            }
            return null;
        }
        public override T Intern(T value)
        {
            if(ReferenceEquals(value, null))
                return null;
            int hash = GetHashCode(value);
            int idx = IndexFromHash(hash, _indices.Length);
            WeakReference yoink = null;
        	for(int i = _indices[idx] - 1; i >= 0; i = _slots[i].Next)
        	{
        	    if(_slots[i].Hash == hash)
        	    {
        	        WeakReference wr = _slots[i].Reference;
        	        T target = (T)wr.Target;
        	        if(ReferenceEquals(target, null))
        	        {
        	            if(yoink == null)
        	                yoink = wr;
        	        }
        	        else if(Equals(target, value))
            		    return target;
        	    }
        	}
        	if(yoink != null)//A matching weak reference going unused.
        	    yoink.Target = value;
        	else
        	{
            	int addIdx;
            	if(_freeList >= 0)
            	{
            		addIdx = _freeList;
            		_freeList = _slots[addIdx].Next;
            	}
            	else
            	{
            		if(_addingIndex == _slots.Length)
            		{
            			IncreaseCapacity();
            			idx = IndexFromHash(hash, _indices.Length);
            		}
            		addIdx = _addingIndex;
            		++_addingIndex;
            	}
            	_slots[addIdx].Hash = hash;
            	_slots[addIdx].Reference = new WeakReference(value);
            	_slots[addIdx].Next = _indices[idx] - 1;
            	_indices[idx] = addIdx + 1;
        	}
        	return value;
        }
        private void IncreaseCapacity()
        {
            if(!ClearDead())
            {
            	int newSize = InternHelper.NextSize(2 * _addingIndex);
            	if(newSize <= _addingIndex)
            	    throw new ArgumentException("Intern Pool Overflow");
            	WeakSlot[] newSlots = new WeakSlot[newSize];
        		Array.Copy(_slots, 0, newSlots, 0, _addingIndex);
            	int[] _newIndices = new int[newSize];
            	for(int i = 0; i < _addingIndex; ++i)
            	{
            	    int idx = IndexFromHash(newSlots[i].Hash, newSize);
            		newSlots[i].Next = _newIndices[idx] - 1;
            		_newIndices[idx] = i + 1;
            	}
            	_slots = newSlots;
            	_indices = _newIndices;
            }
        }
        private bool ClearDead()
        {
            for(int i = 0; i < _addingIndex; ++i)
                if(!_slots[i].Reference.IsAlive)
                {
                    CopyAlive();
                    return true;
                }
            return false;
        }
        private void CopyAlive()
        {
            int len = _indices.Length;
            int[] newIndices = new int[len];
            WeakSlot[] newSlots = new WeakSlot[len];
            int newFreeList = -1;
            int newAddingIndex = 0;
            for(var i = 0; i != _slots.Length; ++i)
            {
                T toCopy = (T)_slots[i].Reference.Target;
                if(!ReferenceEquals(toCopy, null))
                {
                    int hash = _slots[i].Hash;
                    int idx = IndexFromHash(hash, len);
                	int addIdx;
                	if(newFreeList >= 0)
                	{
                		addIdx = newFreeList;
                		newFreeList = newSlots[addIdx].Next;
                	}
                	else
                	{
                		addIdx = newAddingIndex;
                		++newAddingIndex;
                	}
                	newSlots[addIdx].Hash = hash;
                	newSlots[addIdx].Reference = new WeakReference(toCopy);
                	newSlots[addIdx].Next = newIndices[idx] - 1;
                	newIndices[idx] = addIdx + 1;
                }
            }
            _indices = newIndices;
            _slots = newSlots;
            _freeList = newFreeList;
            _addingIndex = newAddingIndex;
        }
    }
    
    public static partial class MoreEnumerable
    {
        /// <summary>
        /// 
        /// </summary>
        /// <param name="source"></param>
        /// <param name="comparer"></param>
        /// <returns></returns>
        public static IEnumerable<T> Intern<T>(this IEnumerable<T> source, IEqualityComparer<T> comparer) where T : class
        {
            return new StrongInternPool<T>(comparer).Intern(source);
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="source"></param>
        /// <returns></returns>
        public static IEnumerable<T> Intern<T>(this IEnumerable<T> source) where T : class
        {
            return new StrongInternPool<T>(EqualityComparer<T>.Default).Intern(source);
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="source"></param>
        /// <param name="pool"></param>
        /// <returns></returns>
        public static IEnumerable<T> Intern<T>(this IEnumerable<T> source, InternPool<T> pool) where T : class
        {
            return pool.Intern(source);
        }
    }
    public sealed class SystemInternPool : InternPool<string> /* wrap the system string interning*/
    {
        public SystemInternPool()
            :base(EqualityComparer<string>.Default)
        {           
        }
        public override string Intern(string value)
        {
            return string.Intern(value);
        }
        public override string IsInterned(string value)
        {
            return string.IsInterned(value);
        }
    }
    public sealed class EmptyInternPool<T> : InternPool<T> where T : class
    {
        public EmptyInternPool()
            :base(EqualityComparer<T>.Default)
        {
        }
        public override T IsInterned(T value)
        {
            return null;
        }
        public override T Intern(T value)
        {
            return value;
        }
    }
}